import hashlib
import pefile
import reque_test
# 计算得到数据长度,自动使用推荐大小
def NumberOfBytesHumanRepresentation(value):
    if value <= 1024:
        return '%s bytes' % value
    elif value < 1024 * 1024:
        return '%.1f KB' % (float(value) / 1024.0)
    elif value < 1024 * 1024 * 1024:
        return '%.1f MB' % (float(value) / 1024.0 / 1024.0)
    else:
        return '%.1f GB' % (float(value) / 1024.0 / 1024.0 / 1024.0)

# 获取PE头部基本信息
def GetHeader(pe, output_file):
    raw = pe.write()

    with open(output_file, 'w', encoding='utf-8') as f:
        # 扫描基本信息
        f.write("-" * 50 + "\n")
        f.write("程序基本信息\n")
        f.write("-" * 50 + "\n")
        if (hex(pe.FILE_HEADER.Machine) == "0x14c"):
            f.write("程序位数: {}\n".format("x86"))
        if (hex(pe.FILE_HEADER.Machine) == "0x8664"):
            f.write("程序位数: {}\n".format("x64"))

        if (hex(pe.FILE_HEADER.Characteristics) == "0x102"):
            f.write("程序类型: Executable\n")
        elif (hex(pe.FILE_HEADER.Characteristics) == "0x2102"):
            f.write("程序类型: Dynamic link library\n")

        if pe.OPTIONAL_HEADER.AddressOfEntryPoint:
            oep = pe.OPTIONAL_HEADER.AddressOfEntryPoint + pe.OPTIONAL_HEADER.ImageBase
            f.write("实际入口: {}\n".format(hex(oep)))

        f.write("映像基址: {}\n".format(hex(pe.OPTIONAL_HEADER.ImageBase)))
        f.write("虚拟入口: {}\n".format(hex(pe.OPTIONAL_HEADER.AddressOfEntryPoint)))
        f.write("映像大小: {}\n".format(hex(pe.OPTIONAL_HEADER.SizeOfImage)))
        f.write("区段对齐: {}\n".format(hex(pe.OPTIONAL_HEADER.SectionAlignment)))
        f.write("文件对齐: {}\n".format(hex(pe.OPTIONAL_HEADER.FileAlignment)))
        f.write("区块数量: {}\n".format(int(pe.FILE_HEADER.NumberOfSections + 1)))
        f.write('熵值比例: %f (Min=0.0, Max=8.0)\n' % pe.sections[0].entropy_H(raw))

        # 计算压缩数据
        f.write("-" * 50 + "\n")
        f.write("计算压缩数据\n")
        f.write("-" * 50 + "\n")
        f.write('MD5     : %s\n' % hashlib.md5(raw).hexdigest())
        f.write('SHA-1   : %s\n' % hashlib.sha1(raw).hexdigest())
        f.write('SHA-256 : %s\n' % hashlib.sha256(raw).hexdigest())
        f.write('SHA-512 : %s\n' % hashlib.sha512(raw).hexdigest())

        # 扫描文件末尾是否存在附加数据
        f.write("-" * 50 + "\n")
        f.write("扫描附加数据\n")
        f.write("-" * 50 + "\n")
        overlayOffset = pe.get_overlay_data_start_offset()
        if overlayOffset is not None:
            f.write("起始文件位置: 0x%08x\n" % overlayOffset)
            overlaySize = len(raw[overlayOffset:])
            f.write("长度: 0x%08x %s %.2f%%\n" % (overlaySize, NumberOfBytesHumanRepresentation(overlaySize), float(overlaySize) / float(len(raw)) * 100.0))
            f.write("MD5: %s\n" % hashlib.md5(raw[overlayOffset:]).hexdigest())
            f.write("SHA-256: %s\n" % hashlib.sha256(raw[overlayOffset:]).hexdigest())
    return str(hashlib.md5(raw).hexdigest())

def gen_trage(filepath):
    pe = pefile.PE(filepath)
    md5 = GetHeader(pe, "basic_info.txt")
    print(md5)
    reque_test.trage_result(md5)

